#######################################################

library(parallel)
library(ComplexHeatmap)
library(optparse)

##########################################################################################
option_list <- list(
    make_option(c("--atac_file"), type = "character"),
    make_option(c("--type"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--cpu"), type = "character"),
    make_option(c("--divide"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 单细胞开放文件
    atac_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/testis_combined_peak.combineRNA.qc.Rdata"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/diff_atac"

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    divide <- 3

    type <- "germ"

    ## cpu
    cpu <- 10
}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

atac_file <- opt$atac_file
cpu <- as.numeric(opt$cpu)
scriptPath <- opt$scriptPath
out_path <- opt$out_path
type <- opt$type
divide <- as.numeric(opt$divide)

dir.create( out_path , recursive = T )

###########################################################################################
## 读数据
a <- load(atac_file)

##########################################################################################
## 已发表文献写好的脚本
source(paste0(scriptPath, "/plotting_config.R"))
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))
source(paste0(scriptPath, "/archr_helpers.R"))
source(paste0(scriptPath, "/GO_wrappers.R"))

###########################################################################################
## 细胞顺序
grp_order2 = c("SSC",
"Differenting&Differented SPG",
"Leptotene",
"Zygotene",
"Patchytene",
"Diplotene",
"Early stage of spermatids",
"Round&ElongateS.tids",
"Sperm",
"Leydig cells",
"Myoid cells",
"Pericytes",
"Sertoli cells",
"Endothelial cells",
"NKT cells",
"Macrophages")

if(type == "germ"){
  grp_order = c("SSC",
  "Differenting&Differented SPG",
  "Leptotene",
  "Zygotene",
  "Patchytene",
  "Diplotene",
  "Early stage of spermatids",
  "Round&ElongateS.tids",
  "Sperm")
}else if(type == "somatic"){
  grp_order = c("Leydig cells",
  "Myoid cells",
  "Pericytes",
  "Sertoli cells",
  "Endothelial cells",
  "NKT cells",
  "Macrophages")
}

SSC_SPG <- c("SSC" , "Differenting&Differented SPG")
SPC <- c("Leptotene" , "Zygotene", "Patchytene",
    "Diplotene" , "Early stage of spermatids")
SPT <- c("Round&ElongateS.tids" , "Sperm")

###########################################################################################
## 提取生殖细胞

if(divide == 3){
  ## 合并细胞为三大类
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SSC_SPG , "SSC_SPG" , testis_combined_peak_combineRNA@cellColData$cell_type)
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SPC , "SPC" , testis_combined_peak_combineRNA@cellColData$cell_type)
  testis_combined_peak_combineRNA@cellColData$cell_type <- ifelse(testis_combined_peak_combineRNA@cellColData$cell_type %in% SPT , "SPT" , testis_combined_peak_combineRNA@cellColData$cell_type)
  grp_order <- c("SSC_SPG" , "SPC" , "SPT")
  ## 不能变为factor，报错Error in cor(estbgdP, obsbgdP) : incompatible dimensions
  #testis_combined_peak_combineRNA@cellColData$cell_type <- factor( testis_combined_peak_combineRNA@cellColData$cell_type , levels = c("SSC_SPG" , "SPC" , "SPT") , order = T )
}else{
  #testis_combined_peak_combineRNA@cellColData$cell_type <- factor( testis_combined_peak_combineRNA@cellColData$cell_type , levels = grp_order[grp_order %in% unique(testis_combined_peak_combineRNA@cellColData$cell_type)] , order = T )
}

germ <- testis_combined_peak_combineRNA

###########################################################################################
## 差异atac
markersGS_germ <- getMarkerFeatures(
    ArchRProj = germ,
    useMatrix = "GeneScoreMatrix", 
    groupBy = "cell_type",
    testMethod = "wilcoxon"
)

## 提取atac矩阵
GSM_se <- getMatrixFromProject(testis_combined_peak_combineRNA, useMatrix="GeneScoreMatrix")
## 所有背景基因
all_genes <- rowData(GSM_se)$name

###########################################################################################
## 提取所有log2fc > 0
markerList_germ <- getMarkers(markersGS_germ, cutOff = "FDR < 1 & Log2FC > 0")
out_file <- paste0( out_path , "/" , "one_vs_other.tsv" )
markerList_germ_out <- data.frame(markerList_germ)
write.table( markerList_germ_out , out_file , row.names = F , quote = F , sep = "\t" )

###########################################################################################
## 富集分析的函数
pathway_enrich <- function(all_genes = all_genes , clust_genes = clust_genes){
  upGO <- rbind(
    calcTopGo(all_genes, interestingGenes=clust_genes, nodeSize=5, ontology="BP") 
    )

  upGO <- upGO[order(as.numeric(upGO$pvalue), decreasing=FALSE),]

  ## 构造GO对象
  geneList <- factor(as.integer(all_genes %in% clust_genes))
  names(geneList) <- all_genes
  GOdata <- suppressMessages(new(
    "topGOdata",
    ontology = "BP",
    allGenes = geneList,
    annot = annFUN.org, mapping = "org.Hs.eg.db", ID = "symbol",
    nodeSize = 5
  ))
  ## 提取每个通路中感兴趣的基因
  gene_in_pathway <- sapply(upGO$GO.ID, function(x)
    {
      genes <- genesInTerm(GOdata, x)
      # myGenes is the queried gene list
      paste0(genes[[1]][genes[[1]] %in% clust_genes] , collapse = "," )
    })
  
  upGO$gene_in_pathway <- gene_in_pathway

  return(upGO)
}


###########################################################################################
## 对不同的foldchange
#for( logfc in c(0.1 , 0.5 , 1 , 1.5 , 2 , 2.5 , 3) ){

  logfc <- 1
  ############################################
  ## 差异表达画热图
  use_t <- paste0( "FDR < 0.05 & Log2FC > ", logfc )

  markerList_germ <- getMarkers(markersGS_germ, cutOff = use_t)
  markerList_germ <- markerList_germ[match(grp_order, names(markerList_germ))]

  #Visualize Markers as a heatmap
  heatmap <- plotMarkerHeatmap(
    seMarker = markersGS_germ[,grp_order], 
    cutOff = use_t, 
    nLabel = 5, # It still seems like there's not actually a way to NOT plot any labels
    nPrint = 5,
    binaryClusterRows = TRUE,
    clusterCols = FALSE,
    transpose = FALSE
  )

  out_file <- paste0( out_path , "/" , "one_vs_other.logfc_" , logfc , ".pdf" ) 
  pdf(out_file , height = 10 , width = 7)
  print(draw(heatmap, heatmap_legend_side="bot", annotation_legend_side="bot"))
  dev.off()

  ############################################
  ## 通路富集
  GOresults <- bind_rows(mclapply(unique(grp_order),function(cellT){
    message(sprintf("Running GO enrichments on %s...", cellT))
    
    ## 提取显著上调的基因
    clust_genes <- subset( markerList_germ_out , group_name==cellT & Log2FC > logfc & FDR < 0.05 )$name
    up_pathway <- pathway_enrich(all_genes = all_genes , clust_genes = clust_genes)
    up_pathway$up_down <- "up"

    tmp <- up_pathway
    tmp$cell_compare <- cellT

    return(tmp)

  },mc.cores=cpu))

  out_file <- paste0( out_path , "/" , "one_vs_other.logfc_" , logfc , ".GOBP.tsv" )
  write.table( GOresults , out_file , row.names = F , quote = F , sep ="\t" )

#}